package org.shj.demo;

import com.hankcs.hanlp.classification.classifiers.IClassifier;
import com.hankcs.hanlp.classification.classifiers.NaiveBayesClassifier;
import com.hankcs.hanlp.classification.models.AbstractModel;
import com.hankcs.hanlp.classification.models.NaiveBayesModel;
import com.hankcs.hanlp.corpus.io.IOUtil;
import org.shj.file.CaseInfo;
import org.shj.utils.DbUtils;

import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.Map;

/**
 * @author Shen Huang Jian
 * @date 2020-09-29 10:40
 */
public class ClassifyDemo {

    public static void main(String[] args){
        String corpusPath = "D:\\test\\hanlp\\xtmainlist";
        String modelPath = "D:\\test\\hanlp\\model\\order.bin";

        try {
            NaiveBayesModel model = trainOrLoadModel(corpusPath, modelPath);
            IClassifier classifier = new NaiveBayesClassifier(model);

            //String sql = "select question_desc from xt_mainlist where (ques_type_value is null or ques_type_value = '') and question_desc like '%挪车%' limit 100;";
            String sql = "select question_desc from xt_mainlist where ques_type_value is null or ques_type_value = '' limit 100;";

            List<CaseInfo> caseInfos = DbUtils.findBySql(sql, CaseInfo.class);

            for(CaseInfo caseInfo : caseInfos){
                predict(classifier, caseInfo.getQuestionDesc());
            }

//            String text = "违法占道广告牌";
//            //String text = "地下车库堆放杂物";
//            predict(classifier, text);

            /*Map<String, Double> predict = classifier.predict(text);
            for(String key : predict.keySet()){
                System.out.println(key + "-->" + predict.get(key));
            }*/
        }catch (Exception e){

        }

    }

    private static void predict(IClassifier classifier, String text){
        System.out.printf("《%s》 属于分类 【%s】\n", text, classifier.classify(text));
    }

    private static NaiveBayesModel trainOrLoadModel(String corpusPath, String modelPath) throws IOException{
        NaiveBayesModel model = (NaiveBayesModel) IOUtil.readObjectFrom(modelPath);
        if (model != null) {
            return model;
        }

        File corpusFolder = new File(corpusPath);
        if(!corpusFolder.exists() || corpusFolder.isFile()){
            System.err.println("分类语料库（" + corpusPath + "）不存在");
            System.exit(1);
        }

        IClassifier classifier = new NaiveBayesClassifier();
        classifier.train(corpusPath);
        model = ((NaiveBayesClassifier) classifier).getNaiveBayesModel();
        IOUtil.saveObjectTo(model, modelPath);
        return model;
    }
}
